{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 4: Input and Output\n",
    "\n",
    "In this tutorial, you'll learn:\n",
    "\n",
    "- Generating input patterns (Poisson, periodic, custom)\n",
    "- Input encoding strategies\n",
    "- Using readout layers\n",
    "- Population coding and decoding\n",
    "- Recording and analyzing network outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainpy\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import braintools\n",
    "import matplotlib.pyplot as plt\n",
    "import jax.numpy as jnp\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Understanding Inputs and Outputs\n",
    "\n",
    "Neural networks need:\n",
    "\n",
    "**Inputs** → Convert external signals to neural activity\n",
    "- Current injection\n",
    "- Spike trains (Poisson, regular)\n",
    "- Temporal patterns\n",
    "\n",
    "**Outputs** → Extract information from network\n",
    "- Spike counts\n",
    "- Population vectors\n",
    "- Readout layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Constant Current Input\n",
    "\n",
    "The simplest input: constant current."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "brainstate.environ.set(dt=0.1 * u.ms)\n",
    "\n",
    "# Create neuron\n",
    "neuron = brainpy.state.LIF(10, V_rest=-65.*u.mV, V_th=-50.*u.mV, tau=10.*u.ms)\n",
    "brainstate.nn.init_all_states(neuron)\n",
    "\n",
    "# Simulate with constant input\n",
    "duration = 200. * u.ms\n",
    "times = u.math.arange(0.*u.ms, duration, brainstate.environ.get_dt())\n",
    "\n",
    "I_constant = 2.0 * u.nA\n",
    "spikes = brainstate.transform.for_loop(\n",
    "    lambda t: neuron(I_constant),\n",
    "    times\n",
    ")\n",
    "\n",
    "# Plot\n",
    "t_idx, n_idx = u.math.where(spikes != 0)\n",
    "plt.figure(figsize=(10, 4))\n",
    "plt.scatter(times[t_idx].to_decimal(u.ms), n_idx, s=5, c='black')\n",
    "plt.xlabel('Time (ms)')\n",
    "plt.ylabel('Neuron Index')\n",
    "plt.title('Response to Constant Current Input')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.show()\n",
    "\n",
    "print(f\"Total spikes: {len(t_idx)}\")\n",
    "print(f\"Average rate: {len(t_idx) / (10 * duration.to_decimal(u.second)):.2f} Hz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Poisson Spike Trains\n",
    "\n",
    "Realistic input: random Poisson spike trains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poisson_input(size, rate, dt):\n",
    "    \"\"\"Generate Poisson spike train.\n",
    "    \n",
    "    Args:\n",
    "        size: Number of neurons\n",
    "        rate: Firing rate (Hz)\n",
    "        dt: Time step\n",
    "    \n",
    "    Returns:\n",
    "        Binary spike array\n",
    "    \"\"\"\n",
    "    prob = rate * dt.to_decimal(u.second)\n",
    "    return (brainstate.random.rand(size) < prob).astype(float)\n",
    "\n",
    "# Test Poisson input\n",
    "brainstate.nn.init_all_states(neuron)\n",
    "rate = 50 * u.Hz\n",
    "dt = brainstate.environ.get_dt()\n",
    "\n",
    "input_spikes_hist = []\n",
    "output_spikes_hist = []\n",
    "\n",
    "for t in times:\n",
    "    # Generate Poisson input\n",
    "    input_spikes = poisson_input(10, rate, dt)\n",
    "    input_spikes_hist.append(input_spikes)\n",
    "    \n",
    "    # Convert spikes to current (simple model)\n",
    "    I_poisson = input_spikes * 5.0 * u.nA\n",
    "    neuron(I_poisson)\n",
    "    output_spikes_hist.append(neuron.get_spike())\n",
    "\n",
    "input_spikes_hist = jnp.array(input_spikes_hist)\n",
    "output_spikes_hist = u.math.asarray(output_spikes_hist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize input and output\n",
    "fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\n",
    "\n",
    "# Input spikes\n",
    "t_in, n_in = jnp.where(input_spikes_hist > 0)\n",
    "axes[0].scatter(times[t_in].to_decimal(u.ms), n_in, s=2, c='blue', alpha=0.5)\n",
    "axes[0].set_ylabel('Neuron Index')\n",
    "axes[0].set_title(f'Input: Poisson Spike Train ({rate.to_decimal(u.Hz):.0f} Hz)', fontweight='bold')\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Output spikes\n",
    "t_out, n_out = u.math.where(output_spikes_hist != 0)\n",
    "axes[1].scatter(times[t_out].to_decimal(u.ms), n_out, s=2, c='red', alpha=0.5)\n",
    "axes[1].set_xlabel('Time (ms)')\n",
    "axes[1].set_ylabel('Neuron Index')\n",
    "axes[1].set_title('Output: Neuron Response', fontweight='bold')\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"Input spikes: {len(t_in)}\")\n",
    "print(f\"Output spikes: {len(t_out)}\")\n",
    "print(f\"Gain: {len(t_out) / len(t_in):.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 4: Periodic Input Patterns\n",
    "\n",
    "Regular, rhythmic inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def periodic_input(t, frequency, amplitude, phase=0):\n",
    "    \"\"\"Generate sinusoidal input current.\n",
    "    \n",
    "    Args:\n",
    "        t: Time\n",
    "        frequency: Oscillation frequency\n",
    "        amplitude: Current amplitude\n",
    "        phase: Phase offset\n",
    "    \"\"\"\n",
    "    omega = 2 * jnp.pi * frequency.to_decimal(u.Hz)\n",
    "    t_sec = t.to_decimal(u.second)\n",
    "    return amplitude * (0.5 + 0.5 * jnp.sin(omega * t_sec + phase))\n",
    "\n",
    "# Test periodic input\n",
    "brainstate.nn.init_all_states(neuron)\n",
    "freq = 10 * u.Hz\n",
    "amp = 3.0 * u.nA\n",
    "\n",
    "currents_hist = []\n",
    "spikes_hist = []\n",
    "\n",
    "for t in times:\n",
    "    I_periodic = periodic_input(t, freq, amp)\n",
    "    currents_hist.append(I_periodic)\n",
    "    neuron(I_periodic)\n",
    "    spikes_hist.append(neuron.get_spike())\n",
    "\n",
    "currents_hist = u.math.asarray(currents_hist)\n",
    "spikes_hist = u.math.asarray(spikes_hist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot periodic input and response\n",
    "fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\n",
    "\n",
    "# Input current\n",
    "axes[0].plot(times.to_decimal(u.ms), currents_hist.to_decimal(u.nA), \n",
    "            linewidth=2, color='blue')\n",
    "axes[0].set_ylabel('Current (nA)')\n",
    "axes[0].set_title(f'Input: Periodic Current ({freq})', fontweight='bold')\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Output spikes\n",
    "t_idx, n_idx = u.math.where(spikes_hist != 0)\n",
    "axes[1].scatter(times[t_idx].to_decimal(u.ms), n_idx, s=5, c='red', alpha=0.7)\n",
    "axes[1].set_xlabel('Time (ms)')\n",
    "axes[1].set_ylabel('Neuron Index')\n",
    "axes[1].set_title('Output: Phase-Locked Spiking', fontweight='bold')\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Observation: Neurons fire preferentially during high-current phases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 5: Rate Coding\n",
    "\n",
    "Encode information in firing rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rate_encode(values, max_rate, dt):\n",
    "    \"\"\"Encode values as Poisson spike trains.\n",
    "    \n",
    "    Args:\n",
    "        values: Array of values to encode (0 to 1)\n",
    "        max_rate: Maximum firing rate\n",
    "        dt: Time step\n",
    "    \n",
    "    Returns:\n",
    "        Binary spike array\n",
    "    \"\"\"\n",
    "    rates = values * max_rate.to_decimal(u.Hz)\n",
    "    probs = rates * dt.to_decimal(u.second)\n",
    "    return (brainstate.random.rand(len(values)) < probs).astype(float)\n",
    "\n",
    "# Example: encode a sine wave\n",
    "n_neurons = 10\n",
    "max_rate = 100 * u.Hz\n",
    "duration = 500. * u.ms\n",
    "times = u.math.arange(0.*u.ms, duration, brainstate.environ.get_dt())\n",
    "\n",
    "encoded_spikes = []\n",
    "signal_values = []\n",
    "\n",
    "for i, t in enumerate(times):\n",
    "    # Signal to encode (sine wave)\n",
    "    signal = 0.5 + 0.5 * jnp.sin(2 * jnp.pi * 5 * t.to_decimal(u.second))\n",
    "    signal_values.append(signal)\n",
    "    \n",
    "    # Encode as spikes for each neuron\n",
    "    values = jnp.ones(n_neurons) * signal  # Same value for all neurons\n",
    "    spikes = rate_encode(values, max_rate, brainstate.environ.get_dt())\n",
    "    encoded_spikes.append(spikes)\n",
    "\n",
    "encoded_spikes = jnp.array(encoded_spikes)\n",
    "signal_values = jnp.array(signal_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize rate coding\n",
    "fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\n",
    "\n",
    "# Original signal\n",
    "axes[0].plot(times.to_decimal(u.ms), signal_values, linewidth=2, color='blue')\n",
    "axes[0].set_ylabel('Signal Value')\n",
    "axes[0].set_title('Original Signal (to be encoded)', fontweight='bold')\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Encoded spikes\n",
    "t_idx, n_idx = jnp.where(encoded_spikes > 0)\n",
    "axes[1].scatter(times[t_idx].to_decimal(u.ms), n_idx, s=1, c='red', alpha=0.5)\n",
    "axes[1].set_xlabel('Time (ms)')\n",
    "axes[1].set_ylabel('Neuron Index')\n",
    "axes[1].set_title('Rate-Coded Spike Train', fontweight='bold')\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Higher signal → higher spike density\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 6: Population Coding\n",
    "\n",
    "Multiple neurons encode a single value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def population_encode(value, n_neurons, pref_values, sigma, max_rate, dt):\n",
    "    \"\"\"Encode value using population code with tuning curves.\n",
    "    \n",
    "    Args:\n",
    "        value: Value to encode (0 to 1)\n",
    "        n_neurons: Number of neurons\n",
    "        pref_values: Preferred values for each neuron\n",
    "        sigma: Tuning width\n",
    "        max_rate: Maximum firing rate\n",
    "        dt: Time step\n",
    "    \"\"\"\n",
    "    # Tuning curves: Gaussian around preferred value\n",
    "    responses = jnp.exp(-0.5 * ((value - pref_values) / sigma)**2)\n",
    "    rates = responses * max_rate.to_decimal(u.Hz)\n",
    "    probs = rates * dt.to_decimal(u.second)\n",
    "    return (brainstate.random.rand(n_neurons) < probs).astype(float)\n",
    "\n",
    "# Setup population\n",
    "n_pop = 20\n",
    "pref_values = jnp.linspace(0, 1, n_pop)  # Evenly spaced preferences\n",
    "sigma = 0.2\n",
    "max_rate = 100 * u.Hz\n",
    "\n",
    "# Encode a slowly changing value\n",
    "duration = 500. * u.ms\n",
    "times = u.math.arange(0.*u.ms, duration, brainstate.environ.get_dt())\n",
    "\n",
    "pop_spikes = []\n",
    "true_values = []\n",
    "\n",
    "for i, t in enumerate(times):\n",
    "    # Value changes over time\n",
    "    value = 0.5 + 0.3 * jnp.sin(2 * jnp.pi * 2 * t.to_decimal(u.second))\n",
    "    true_values.append(value)\n",
    "    \n",
    "    # Population encoding\n",
    "    spikes = population_encode(value, n_pop, pref_values, sigma, max_rate, \n",
    "                               brainstate.environ.get_dt())\n",
    "    pop_spikes.append(spikes)\n",
    "\n",
    "pop_spikes = jnp.array(pop_spikes)\n",
    "true_values = jnp.array(true_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize population coding\n",
    "fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\n",
    "\n",
    "# True value\n",
    "axes[0].plot(times.to_decimal(u.ms), true_values, linewidth=2, color='blue')\n",
    "axes[0].set_ylabel('Encoded Value')\n",
    "axes[0].set_title('True Value (to be encoded)', fontweight='bold')\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Population spikes\n",
    "t_idx, n_idx = jnp.where(pop_spikes > 0)\n",
    "axes[1].scatter(times[t_idx].to_decimal(u.ms), n_idx, s=2, c='red', alpha=0.5)\n",
    "axes[1].set_xlabel('Time (ms)')\n",
    "axes[1].set_ylabel('Neuron Index (Preference)')\n",
    "axes[1].set_title('Population Code: Activity Follows Value', fontweight='bold')\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Peak activity shifts with encoded value\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 7: Population Decoding\n",
    "\n",
    "Extract the encoded value from population activity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def population_decode(spike_counts, pref_values):\n",
    "    \"\"\"Decode value from population activity.\n",
    "    \n",
    "    Args:\n",
    "        spike_counts: Number of spikes per neuron\n",
    "        pref_values: Preferred values of neurons\n",
    "    \n",
    "    Returns:\n",
    "        Decoded value (population vector)\n",
    "    \"\"\"\n",
    "    # Population vector: weighted average\n",
    "    total_activity = jnp.sum(spike_counts)\n",
    "    if total_activity > 0:\n",
    "        decoded = jnp.sum(spike_counts * pref_values) / total_activity\n",
    "        return decoded\n",
    "    else:\n",
    "        return 0.5  # Default\n",
    "\n",
    "# Decode the population activity\n",
    "window_size = 50  # ms\n",
    "window_steps = int(window_size / brainstate.environ.get_dt().to_decimal(u.ms))\n",
    "\n",
    "decoded_values = []\n",
    "decode_times = []\n",
    "\n",
    "for i in range(0, len(times) - window_steps, window_steps // 2):\n",
    "    # Count spikes in window\n",
    "    window_spikes = pop_spikes[i:i+window_steps]\n",
    "    spike_counts = jnp.sum(window_spikes, axis=0)\n",
    "    \n",
    "    # Decode\n",
    "    decoded = population_decode(spike_counts, pref_values)\n",
    "    decoded_values.append(decoded)\n",
    "    decode_times.append(times[i + window_steps//2])\n",
    "\n",
    "decoded_values = jnp.array(decoded_values)\n",
    "decode_times = u.math.asarray(decode_times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare true and decoded values\n",
    "plt.figure(figsize=(12, 5))\n",
    "plt.plot(times.to_decimal(u.ms), true_values, linewidth=2, \n",
    "        label='True Value', color='blue', alpha=0.7)\n",
    "plt.plot(decode_times.to_decimal(u.ms), decoded_values, linewidth=2, \n",
    "        label='Decoded Value', color='red', linestyle='--', alpha=0.7)\n",
    "plt.xlabel('Time (ms)', fontsize=12)\n",
    "plt.ylabel('Value', fontsize=12)\n",
    "plt.title('Population Decoding: True vs Decoded Values', fontsize=14, fontweight='bold')\n",
    "plt.legend(fontsize=11)\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Calculate decoding error\n",
    "# Interpolate true values at decode times\n",
    "true_at_decode = jnp.interp(\n",
    "    decode_times.to_decimal(u.ms),\n",
    "    times.to_decimal(u.ms),\n",
    "    true_values\n",
    ")\n",
    "error = jnp.abs(decoded_values - true_at_decode)\n",
    "print(f\"Mean decoding error: {jnp.mean(error):.4f}\")\n",
    "print(f\"Max decoding error: {jnp.max(error):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 8: Readout Layers\n",
    "\n",
    "Use a readout layer to extract output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create network with readout\n",
    "class NetworkWithReadout(brainstate.nn.Module):\n",
    "    def __init__(self, n_input, n_hidden, n_output):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Hidden layer (recurrent LIF neurons)\n",
    "        self.hidden = brainpy.state.LIF(n_hidden, V_rest=-65.*u.mV, V_th=-50.*u.mV, tau=10.*u.ms)\n",
    "        \n",
    "        # Readout layer\n",
    "        self.readout = brainpy.state.Readout(\n",
    "            n_hidden, n_output,\n",
    "            weight_initializer=braintools.init.KaimingNormal()\n",
    "        )\n",
    "\n",
    "    \n",
    "    def update(self, spike_input):\n",
    "        # Convert input spikes to current\n",
    "        I_input = spike_input * 5.0 * u.nA\n",
    "        \n",
    "        # Update hidden neurons\n",
    "        self.hidden(I_input)\n",
    "        spikes = self.hidden.get_spike()\n",
    "        \n",
    "        # Readout\n",
    "        output = self.readout(spikes)\n",
    "        \n",
    "        return output, spikes\n",
    "\n",
    "# Create network\n",
    "net = NetworkWithReadout(n_input=10, n_hidden=50, n_output=2)\n",
    "brainstate.nn.init_all_states(net)\n",
    "\n",
    "print(\"Network with readout layer created\")\n",
    "print(f\"Hidden neurons: {50}\")\n",
    "print(f\"Output dimensions: {2}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test readout\n",
    "duration = 200. * u.ms\n",
    "times = u.math.arange(0.*u.ms, duration, brainstate.environ.get_dt())\n",
    "\n",
    "outputs_hist = []\n",
    "spikes_hist = []\n",
    "\n",
    "for t in times:\n",
    "    # Generate Poisson input\n",
    "    input_spikes = poisson_input(10, 50*u.Hz, brainstate.environ.get_dt())\n",
    "    \n",
    "    # Network update\n",
    "    output, spikes = net.update(input_spikes)\n",
    "    outputs_hist.append(output)\n",
    "    spikes_hist.append(spikes)\n",
    "\n",
    "outputs_hist = jnp.array(outputs_hist)\n",
    "spikes_hist = u.math.asarray(spikes_hist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize readout\n",
    "fig, axes = plt.subplots(2, 1, figsize=(12, 8), sharex=True)\n",
    "\n",
    "# Hidden layer activity\n",
    "t_idx, n_idx = u.math.where(spikes_hist != 0)\n",
    "axes[0].scatter(times[t_idx].to_decimal(u.ms), n_idx, s=1, c='blue', alpha=0.5)\n",
    "axes[0].set_ylabel('Neuron Index')\n",
    "axes[0].set_title('Hidden Layer Spikes', fontweight='bold')\n",
    "axes[0].grid(True, alpha=0.3)\n",
    "\n",
    "# Readout outputs\n",
    "axes[1].plot(times.to_decimal(u.ms), outputs_hist[:, 0], \n",
    "            linewidth=2, label='Output 1', alpha=0.7)\n",
    "axes[1].plot(times.to_decimal(u.ms), outputs_hist[:, 1], \n",
    "            linewidth=2, label='Output 2', alpha=0.7)\n",
    "axes[1].set_xlabel('Time (ms)')\n",
    "axes[1].set_ylabel('Readout Value')\n",
    "axes[1].set_title('Readout Layer Outputs', fontweight='bold')\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"Readout layer converts spikes to continuous values\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 9: Recording Network States\n",
    "\n",
    "Record variables during simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manual recording example\n",
    "neuron = brainpy.state.LIF(5, V_rest=-65.*u.mV, V_th=-50.*u.mV, tau=10.*u.ms)\n",
    "brainstate.nn.init_all_states(neuron)\n",
    "\n",
    "duration = 100. * u.ms\n",
    "times = u.math.arange(0.*u.ms, duration, brainstate.environ.get_dt())\n",
    "\n",
    "# Preallocate recording arrays\n",
    "n_steps = len(times)\n",
    "V_hist = []\n",
    "spike_hist = []\n",
    "\n",
    "for t in times:\n",
    "    neuron(2.0 * u.nA)\n",
    "    \n",
    "    # Record states\n",
    "    V_hist.append(neuron.V.value.copy())\n",
    "    spike_hist.append(neuron.get_spike().copy())\n",
    "\n",
    "V_hist = u.math.asarray(V_hist)  # Shape: (time, neurons)\n",
    "spike_hist = u.math.asarray(spike_hist)\n",
    "\n",
    "print(f\"Recorded {n_steps} time steps\")\n",
    "print(f\"Voltage history shape: {V_hist.shape}\")\n",
    "print(f\"Spike history shape: {spike_hist.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot recorded states\n",
    "plt.figure(figsize=(12, 6))\n",
    "\n",
    "# Plot voltage traces\n",
    "for i in range(5):\n",
    "    V_trace = V_hist[:, i]\n",
    "    # Mark spikes\n",
    "    spike_times = times[spike_hist[:, i] > 0]\n",
    "    V_with_spikes = V_trace.copy()\n",
    "    V_with_spikes = V_with_spikes.to_decimal(u.mV)\n",
    "    \n",
    "    plt.plot(times.to_decimal(u.ms), V_with_spikes, \n",
    "            linewidth=1.5, alpha=0.7, label=f'Neuron {i}')\n",
    "\n",
    "plt.axhline(y=-50, color='r', linestyle='--', alpha=0.5, label='Threshold')\n",
    "plt.xlabel('Time (ms)', fontsize=12)\n",
    "plt.ylabel('Membrane Potential (mV)', fontsize=12)\n",
    "plt.title('Recorded Voltage Traces', fontsize=14, fontweight='bold')\n",
    "plt.legend(loc='upper right', fontsize=9)\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this tutorial, you learned:\n",
    "\n",
    "✅ **Input Generation**: Constant, Poisson, periodic patterns\n",
    "\n",
    "✅ **Rate Coding**: Encoding values as firing rates\n",
    "\n",
    "✅ **Population Coding**: Multiple neurons encode single values\n",
    "\n",
    "✅ **Population Decoding**: Extract values from spike trains\n",
    "\n",
    "✅ **Readout Layers**: Convert spikes to continuous outputs\n",
    "\n",
    "✅ **Recording States**: Track network variables over time\n",
    "\n",
    "## Key Concepts\n",
    "\n",
    "1. **Input Encoding**: Convert signals → spike patterns\n",
    "2. **Population Codes**: Distributed representation across neurons\n",
    "3. **Decoding**: Extract information from population activity\n",
    "4. **Readout**: Linear combination of spike counts\n",
    "\n",
    "## Next Steps\n",
    "\n",
    "- **Tutorial 5**: Learn [SNN training](../advanced/05-snn-training.ipynb)\n",
    "- **Examples**: See [trained networks](../../examples/gallery.rst#snn-training)\n",
    "- **Advanced**: Explore [reservoir computing](../../examples/gallery.rst)\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Implement temporal coding (first-spike latency)\n",
    "2. Create a 2D population code (e.g., for position)\n",
    "3. Build a classifier using readout layer\n",
    "4. Compare different decoding methods (vector, maximum likelihood)\n",
    "5. Implement sparse coding with inhibition"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Ecosystem-py",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
